# -*- coding: utf-8 -*-#

# -------------------------------------------------------------------------------
# Name:         zx
# Description:
# Author:       zx
# Date:         2021/4/1
# -------------------------------------------------------------------------------
import csv
import os
import random

import numpy as np
from sklearn import tree

FilePath = "/train_data/"
files = os.listdir(os.path.dirname(os.path.abspath(__file__)) + FilePath)  # 得到文件夹下的所有文件名称

# samples = {
#     type : [n_samples, [n_features, [mean, variance]]]
# }

all_samples = {
    "pos": [],
    "neg": [],
}

for file in files:
    # 读取csv至字典
    csvFile = open(os.path.dirname(os.path.abspath(__file__)) + FilePath + file, "r", encoding="GBK")
    reader = csv.reader(csvFile)
    sample = np.array(list(reader)[1:]).astype(np.float32)
    if int(file[6:8]) >= 22:
        all_samples["pos"].append(
            [[np.nanmean(sample[:, i]), np.nanvar(sample[:, i])] for i in range(1, sample.shape[1])]
        )
    else:
        all_samples["neg"].append(
            [[np.nanmean(sample[:, i]), np.nanvar(sample[:, i])] for i in range(1, sample.shape[1])]
        )
    csvFile.close()

all_samples["pos"] = [[_ for item in sample for _ in item] for sample in all_samples["pos"]]
all_samples["neg"] = [[_ for item in sample for _ in item] for sample in all_samples["neg"]]

print(all_samples["pos"])
print(all_samples["neg"])
random.shuffle(all_samples["pos"])
random.shuffle(all_samples["neg"])
print(all_samples["pos"])
print(all_samples["neg"])

pos_N = len(all_samples["pos"])
neg_N = len(all_samples["neg"])
print(pos_N, neg_N)

# 切分训练集和测试集
# 负：16  正：27
neg_index = 15
pos_index = 26
train_data = all_samples["neg"][0:neg_index] + all_samples["pos"][0:pos_index]
train_label = [0] * neg_index + [1] * pos_index

test_data = all_samples["neg"][neg_index:] + all_samples["pos"][pos_index:]
test_label = [0] * (neg_N - neg_index) + [1] * (pos_N - pos_index)
print(len(train_label), len(test_label))

# 再次打乱训练集
index = list(range(len(train_data)))
print(index)
random.shuffle(index)
print(index)
train_data = [train_data[i] for i in index]
train_label = [train_label[i] for i in index]

# 决策树
clf = tree.DecisionTreeRegressor()
clf = clf.fit(train_data, train_label)
predict = clf.predict(test_data)

# 转换
predict = np.array(predict).astype(np.int32)
test_label = np.array(test_label).astype(np.int32)

# 结果
result = predict == test_label
print(predict)
print(test_label)
print(sum(result) / len(result))


def predict(data):
    # 求出每列均值、方差
    data_mean_var = [[np.nanmean(data[:, i]), np.nanvar(data[:, i])] for i in range(1, data.shape[1])]
    # 展开每列均值、方差
    my_data = [_ for item_mean_var in data_mean_var for _ in item_mean_var]
    return int(clf.predict([my_data])[0])
